import torch
import torch.nn as nn
import pandas as pd
import matplotlib.pyplot as plt
import random
import numpy

def generate_real():
    real_data = torch.FloatTensor(
        [random.uniform(0.8,1.0),
         random.uniform(0.0,0.2),
         random.uniform(0.8,1.0),
         random.uniform(0.0,0.2),])
    return real_data

def generate_random():
    random_data = torch.rand(4)
    return random_data

generate_real()

generate_random()

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(4,3),
            nn.Sigmoid(),
            nn.Linear(3,1),
            nn.Sigmoid()
        )
        #self.model = self.model.cuda()
        self.loss_function = nn.MSELoss()
        self.optimiser = torch.optim.SGD(self.parameters(), lr=0.01)
        self.counter = 0
        self.progress = []
        pass
    def forward(self, inputs):
        return self.model(inputs)
    def train(self, inputs, targets):
        outputs = self.forward(inputs)
        loss = self.loss_function(outputs, targets)
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()
        self.counter +=1
        if (self.counter % 10 == 0):
            self.progress.append(loss.item())
            pass
        if (self.counter % 10000 == 0):
            print("counter = ",self.counter, "loss = ",loss.item())
            pass
    def plot_progress(self):
        df = pd.DataFrame(self.progress, columns = ['loss'])
        df.plot(ylim=(0, 1.0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0,0.25,0.5))
        pass

D = Discriminator()
for i in range(10000):
    D.train(generate_real().reshape(1,4), torch.FloatTensor([1.0]).reshape(1,1))
    D.train(generate_random().reshape(1,4), torch.FloatTensor([0.0]).reshape(1,1))
    pass

D.plot_progress()
D.forward(generate_real()).item()
D.forward(generate_random()).item()
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(1,3),
            nn.Sigmoid(),
            nn.Linear(3,4),
            nn.Sigmoid()
        )
        #self.model = self.model.cuda()
        self.optimiser = torch.optim.SGD(self.parameters(), lr=0.01)
        self.counter = 0
        self.progress = []
        pass
    def forward(self, inputs):
        return self.model(inputs)
    def train(self, D, inputs, targets):
        g_outputs = self.forward(inputs)
        d_outputs = D.forward(g_outputs)
        loss = D.loss_function(d_outputs, targets)
        self.optimiser.zero_grad()
        loss.backward()
        self.optimiser.step()
        self.counter +=1
        if (self.counter % 10 == 0):
            self.progress.append(loss.item())
            pass
        if (self.counter % 10000 == 0):
            print("counter = ",self.counter, "loss = ",loss.item())
            pass
    def plot_progress(self):
        df = pd.DataFrame(self.progress, columns = ['loss'])
        df.plot(ylim=(0, 1.0), figsize=(16,8), alpha=0.1, marker='.', grid=True, yticks=(0,0.25,0.5))
        pass

G = Generator()
G.forward(torch.FloatTensor([0.5]))

#训练GAN
D = Discriminator()
G = Generator()
image_list = []
for i in range(10000):
    D.train(generate_real(), torch.FloatTensor([1.0]))
    D.train(G.forward(torch.FloatTensor([0.5])).detach(), torch.FloatTensor([0.0]))
    G.train(D, torch.FloatTensor([0.5]), torch.FloatTensor([1.0]))
    if(i % 1000 == 0):
        image_list.append(G.forward(torch.FloatTensor([0.5])).detach().numpy())
    pass


D.plot_progress()
G.plot_progress()
G.forward(torch.FloatTensor([0.5])).detach()
plt.figure(figsize = (16,8))
plt.imshow(numpy.array(image_list).T, interpolation='none', cmap='Blues')